import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import SimpleITK as sitk
import numpy as np
import os
from glob import glob
from matplotlib.colors import ListedColormap,  LinearSegmentedColormap, BoundaryNorm
import matplotlib

matplotlib.use('Agg')

my_cmap = ListedColormap(['black', 'green', 'red', 'yellow'])
bounds = [0.0, 0.5, 1.5, 2.5, 3.5]
norm = BoundaryNorm(bounds, my_cmap.N)


def load_nii(path):
    image = sitk.ReadImage(path)
    image = sitk.GetArrayFromImage(image)
    return image


def trans_brats_label(x):
    mask_WT = x.copy()
    mask_WT[mask_WT == 1] = 1
    mask_WT[mask_WT == 2] = 1
    mask_WT[mask_WT == 3] = 1

    mask_TC = x.copy()
    mask_TC[mask_TC == 1] = 1
    mask_TC[mask_TC == 2] = 0
    mask_TC[mask_TC == 3] = 1

    mask_ET = x.copy()
    mask_ET[mask_ET == 1] = 0
    mask_ET[mask_ET == 2] = 0
    mask_ET[mask_ET == 3] = 1
    
    mask = np.stack([mask_WT, mask_TC, mask_ET], axis=0)
    return mask


def get_image(path):
    image_dic = {
        't1n': None,
        't2w': None,
        't1c': None,
        't2f': None,
        'seg': None
    }
    path = glob(path + '/*')

    for img_path in path:
        if 'seg' in img_path:
            tmp = load_nii(img_path)
            image_dic['seg'] = trans_brats_label(tmp)       
        elif 't1c' in img_path:
            image_dic['t1c'] = load_nii(img_path)
        elif 't1n' in img_path:
            image_dic['t1n'] = load_nii(img_path)
        elif 't2f' in img_path:
            image_dic['t2f'] = load_nii(img_path)
        elif 't2w' in img_path:
            image_dic['t2w'] = load_nii(img_path)
            
    return image_dic


def visual_brats(path, out_dir):
    """_summary_
        可视化四个模态的图像，每个类别原始标签，以及 BraTS 的子区域

        Args:
            path_list (dic): {'image': [t1, t2, t1c, t2f],
                              'label':[]
                              }
    """
    data = get_image(path)

    for i in range(155):
        plt.ioff()
        fig = plt.figure(figsize=(20, 10))
        gs = gridspec.GridSpec(nrows=2, ncols=4, height_ratios=[1, 1])
        
        # the top row is four modality images
        ax0 = fig.add_subplot(gs[0, 0])
        ax0.imshow(data['t2f'][i], cmap='bone')
        ax0.set_title("FLAIR", fontsize=18, weight='bold', y=-0.2)
        ax0.set_xticks([])
        ax0.set_yticks([])

        ax1 = fig.add_subplot(gs[0, 1])
        ax1.imshow(data['t1n'][i], cmap='bone')
        ax1.set_title("T1", fontsize=18, weight='bold', y=-0.2)
        ax1.set_xticks([])
        ax1.set_yticks([])

        ax2 = fig.add_subplot(gs[0, 2])
        ax2.imshow(data['t2w'][i], cmap='bone')
        ax2.set_title("T2", fontsize=18, weight='bold', y=-0.2)
        ax2.set_xticks([])
        ax2.set_yticks([])

        ax3 = fig.add_subplot(gs[0, 3])
        ax3.imshow(data['t1c'][i], cmap='bone')
        ax3.set_title("T1 contrast", fontsize=18, weight='bold', y=-0.2)
        ax3.set_xticks([])
        ax3.set_yticks([])
    
        # the middle row is three subregions of label
        ax10 = fig.add_subplot(gs[1, 0])
        ax10.imshow(data['seg'][0, i], cmap='gray')
        ax10.set_title("WT", fontsize=18, weight='bold', y=-0.2)
        ax10.set_xticks([])
        ax10.set_yticks([])

        ax11 = fig.add_subplot(gs[1, 1])
        ax11.imshow(data['seg'][1, i], cmap='gray')
        ax11.set_title("TC", fontsize=18, weight='bold', y=-0.2)
        ax11.set_xticks([])
        ax11.set_yticks([])

        ax12 = fig.add_subplot(gs[1, 2])
        ax12.imshow(data['seg'][2, i], cmap='gray')
        ax12.set_title("ET", fontsize=18, weight='bold', y=-0.2)
        ax12.set_xticks([])
        ax12.set_yticks([])

        ax13 = fig.add_subplot(gs[1, 3])
        ax13.imshow(np.sum(data['seg'], axis=0)[i], 
                    cmap=my_cmap, 
                    norm=norm)
        ax13.set_title("Labels", fontsize=18, weight='bold', y=-0.2)
        ax13.set_xticks([])
        ax13.set_yticks([])

        plt.suptitle(f"Slice {i}", fontsize=20, weight='bold')

        image_name = os.path.split(path)[-1]
        
        if not os.path.exists(os.path.join(out_dir, image_name)):
            os.mkdir(os.path.join(out_dir, image_name))
            
        des = os.path.join(out_dir, image_name, f"{i}.png")
        fig.savefig(des, 
                    dpi=600,
                    format="png",  
                    pad_inches=0.2, 
                    transparent=False, 
                    bbox_inches='tight')
        print("Case, slice", os.path.split(path)[-1], i)
        plt.clf()
        plt.close(fig)

if __name__ == "__main__":
    path = [
        '/Users/qlc/Desktop/Dataset/Brats2023/Adult_Glioma/TrainingData/BraTS-GLI-00192-000',
        '/Users/qlc/Desktop/Dataset/Brats2023/Adult_Glioma/TrainingData/BraTS-GLI-00192-000/BraTS-GLI-00192-000-seg.nii.gz',
        '/Users/qlc/Desktop/Dataset/Brats2023/Adult_Glioma/TrainingData/BraTS-GLI-00192-000/BraTS-GLI-00192-000-t1c.nii.gz',
        '/Users/qlc/Desktop/Dataset/Brats2023/Adult_Glioma/TrainingData/BraTS-GLI-00192-000/BraTS-GLI-00192-000-t2f.nii.gz',
        '/Users/qlc/Desktop/Dataset/Brats2023/Adult_Glioma/TrainingData/BraTS-GLI-00192-000/BraTS-GLI-00192-000-t1n.nii.gz',
        '/Users/qlc/Desktop/Dataset/Brats2023/Adult_Glioma/TrainingData/BraTS-GLI-00192-000/BraTS-GLI-00192-000-t2w.nii.gz',
    ]
    
    des = '/Users/qlc/Desktop/i.png'
    
    # image = load_nii(path[0])
    # 52  59  47  
    image = get_image(path[0])
    image = image['t2f']
    print(image.shape)
    # print(image.keys())
    
    # image = trans_brats_label(image)
    # image = np.sum(image, axis=0)
    # assert 1 == 2 

    t1c = (image[49, :, :] + image[52, :, :]) / 2

    # seg = image[:, 53, :, :] + image[:, 51, :, :]
    # # seg = image['seg'][:, 53, :, :]
    # seg[seg>0] = 1
    # seg = np.sum(seg, 0)
    # print(np.unique(seg))
    
    fig = plt.figure()
    plt.imshow(t1c, cmap='gray')
    # seg = seg[80:200, 20:140]
    # plt.imshow(seg, cmap=my_cmap, norm=norm)
    plt.xticks([])
    plt.yticks([])
    fig.savefig(des, dpi=600, format="png", pad_inches=0.2, transparent=False, bbox_inches='tight')
